ggsave <- ggplot2::ggsave
body(ggsave)[[11]] <- substitute(invisible())
environment(ggsave) <- asNamespace('ggplot2')
assignInNamespace('ggsave', ggsave, ns = 'ggplot2')
rm(ggsave)


# PPMF handling -----
#' Add prefix to string
#'
#' @param pref string prefix to add
#' @param x string to add the prefix to
#'
#' @return string
add_pref <- function(pref, x){paste0(pref, '_', x)}

#' Aggregate block pop/vap to precinct level
#'
#' This is a less general version of geomander::block2prec() for ppmf data
#'
#' @param ppmf tibble of ppmf data
#' @param prec unquoted
#'
#' @return tibble aggregated to precinct level
ppmf_block2prec <- function(ppmf, prec){
  ppmf %>%
    st_drop_geometry() %>%
    group_by({{prec}}) %>%
    summarize(across(.cols = contains(c('pop', 'vap')), .fns = ~sum(.x, na.rm = TRUE)))
}

# Visualization helpers ----

#' data sourcelabelling
lbl_source = function(x) {
  case_when(x == "das12" ~ "DAS-12.2",
            x == "das19" ~ "DAS-19.61",
            x == "das04" ~ "DAS-4.5",
            x == "das4" ~ "DAS-4.5",
            x == "orig" ~ "Census 2010",
            TRUE ~ "Census 2010")
}
# Colors
# Taken via https://github.com/CoryMcCartan/wacolors/blob/main/R/colors.R
rainier = c(
  lake = "#465177", ragwort = "#E4C22B", lodge = "#965127",
  trees = "#29483A", ground = "#759C44", winter_sky = "#9FB6DA",
  paintbrush = "#DF3383")
skagit = c(red = "#B51829", yellow = "#EFC519", violet = "#831285",
           orange = "#CC7810", purple = "#C886E8", mountains = "#3E6E94")

# DAS palette
PAL_DAS <- c(rainier[c("winter_sky", "ground", "trees")], "#465177")
names(PAL_DAS) <- c("Census 2010", "DAS-12.2", "DAS-4.5", "DAS-19.61")

# Race palette
PAL_race <- skagit[c("red", "yellow", "violet", "orange")]
names(PAL_race) <- c("White", "Black", "Hispanic", "Other")

rm(rainier, skagit)


# Standard themes
theme_ppmf <- function(...) {
  ggplot2::theme_bw(..., base_family="Times", base_size=9) +
    theme(plot.title=element_text(size=9))
}
theme_ppmf_map <- ggplot2::theme_void
palette_ppmf <- scales::manual_pal(values = viridisLite::magma(n = 10))




# Scatter plot ------
#' Compute and format common error statistics
#' @param truth,estimate two numeric vectors of the same length
#' @return A string showing errors
#' @source functions adapted from kuriwaki/ccesMRPrun
#' @examples
#'  error_lbl(1:10, 10:1)
error_lbl <- function(truth, estimate,
                      show_metrics = c("rmse", "mean", "bias"),
                      metrics_lbl = c(rmse = "RMSE", mean = "Mean Abs. Dev", bias = "Mean Dev.")) {

  rmse_stat <- sqrt(mean((truth - estimate)^2))
  mean_stat <- mean(abs(truth - estimate))
  bias_stat <- mean(truth - estimate)

  stat_vec <- c(rmse = rmse_stat, mean = mean_stat, bias = bias_stat)

  show_stat <- paste0(round(stat_vec[show_metrics], 0), " people")

  show_lbl <- str_c(str_c(metrics_lbl[show_metrics], show_stat, sep = ": "),
                    collapse = "\n")

  show_lbl
}


#' Scatter plot with 45 degree line and stats
#'
#' @param xvar Variable for X axis, unquoted
#' @param yvar Variable for Y axis, unquoted
#' @param data The table which includes those variables
#' @param title A string to show for title
#'
#' @source functions adapted from kuriwaki/ccesMRPrun
scatter_counts <- function(xvar, yvar, xlab = NULL, ylab = NULL, title = NULL, data = sc_map) {

  xvar <- enquo(xvar)
  yvar <- enquo(yvar)


  error_txt <- error_lbl(truth = pull(data, !!xvar),
                         estimate = pull(data, !!yvar))


  gg <- ggplot(data) +
    aes(x = {{xvar}}, y = {{yvar}}) +
    geom_point(size = 0.01, alpha = 0.5) +
    geom_abline(alpha = 0.1) +
    theme_classic() +
    scale_x_continuous(trans = scales::pseudo_log_trans(base = 10)) +
    scale_y_continuous(trans = scales::pseudo_log_trans(base = 10)) +
    coord_equal() +
    labs(title = title,
         caption = error_txt) +
    theme(plot.title = element_text(hjust = 0.5, face = "bold", size = 10),
          plot.caption = element_text(color = "darkgray", size = 8),
          axis.title = element_text(size = 9))

  if (!is.null(xlab))
    gg <- gg + labs(x = xlab)
  if (!is.null(ylab))
    gg <- gg + labs(y = ylab)

  gg
}

error_lbl_par <- error_lbl
body(error_lbl_par)[[6]] <- substitute(show_stat <- paste0(round(100*stat_vec[show_metrics], 2), " %"))

#' Create a parity scatterplot
#'
#' @param tb_par tibble of parities
#' @param c2010 true census data
#' @param e DAS data
#' @param lbl label to add to y axis. Default is 'e4-5'
#' @param limx limits for x axis
#' @param limy limits for y axis
#' @param fp return false positive when TRUE and false negative when FALSE
#' @param title
#'
#' @return
scatter_par <- function(tb_par, c2010, e, lbl = 'e4-5',
                        limx = c(.002,.008),
                        limy = c(.002, .008), fp = TRUE,
                        title = 'set title'){

  etxt <- error_lbl_par(rlang::eval_tidy(rlang::enquo(c2010),
                                         tb_par),
                        rlang::eval_tidy(rlang::enquo(e),
                                         tb_par))
  if(!is.null(fp)){
    ftxt <- fnfp(tb_par, !!rlang::enquo(c2010),
                 !!rlang::enquo(e), fp)
  } else {
    ftxt <- ''
  }

  tb_par %>% slice(-1) %>%
    ggplot(aes(x = {{c2010}})) +
    geom_point(aes(y = {{e}})) +
    theme_ppmf() +
    labs(y = paste(lbl, 'Deviation'),
         x = '2010 Census Deviation',
         caption = etxt,
         title = title) +
    coord_fixed() +
    geom_abline(slope = 1, intercept = 0, color = palette_ppmf(5)[5]) +
    geom_vline(xintercept = rlang::eval_tidy(rlang::enquo(c2010), tb_par)[1],
               color = palette_ppmf(10)[8]) +
    geom_hline(yintercept = rlang::eval_tidy(rlang::enquo(c2010), tb_par)[1],
               color = palette_ppmf(10)[8]) +
    theme(plot.caption = element_text(hjust = 1)) +
    scale_y_continuous(limits = limy, labels = scales::percent) +
    scale_x_continuous(limits = limx, labels = scales::percent) +
    annotate('text', x = ifelse(fp, 0.0065, 0.0025),
             y = ifelse(fp, 0.0035, 0.007), label = ftxt, hjust = 0) %>%
    suppressWarnings()
}


# False Negatives ------
#' Get False Negative/Positive Text
#'
#' @param tb_par tibble of parities
#' @param c2010 true census data
#' @param e DAS data
#' @param fp TRUE if false positive or FALSE to report false negative
#'
#' @return string
fnfp <- function(tb_par, c2010, e, fp = TRUE){
  truth <- rlang::eval_tidy(rlang::enquo(c2010), tb_par)
  bound <- truth[1]
  truth <- truth[-1]
  est <- rlang::eval_tidy(rlang::enquo(e), tb_par)[-1]

  drops <- which(truth > bound & est > bound)
  truth <- truth[-drops]
  est <- est[-drops]
  denom <- length(truth)

  if(fp){
    lbl <- paste0('False Acceptance:\n',
                  round(100*(sum(est < bound & truth > bound)/denom), 2), '%')
  } else {
    lbl <- paste0('False Rejection:\n',
                  round(100*(sum(est > bound & truth < bound)/denom),2), '%')
  }
}


# redist diagnostics ------

#' Create a histogram of a VI metrics
#'
#'
#'
#' @param obj A plans_redist object
#' @param ind The index of which draws to use from the `obj`. Currently picks a
#'  sensible sample of 100 draws.
#' @param title Title for the graph. Pick a informative title
#'
#' @param map A map object which has a "pop" column for total population
#'
#' @return A ggplot historgram
hist_VIdist <- function(obj,
                        title,
                        use_ind  = seq(100, n_distinct(obj$draw), length = 100),
                        map = sc_map_orig) {
  samp_plans <- get_plans_matrix(obj)[, use_ind]
  dists_mat <- redist.distances(samp_plans,
                                measure = "info",
                                total_pop = map$pop)[[1]]

  dist_vec <- dists_mat[!diag(dists_mat)]

  tibble(d = dist_vec) %>%
    ggplot(aes(x = d, y = stat(width*density))) +
    geom_histogram(color = "white", bins = 20) +
    scale_y_continuous(expand = expansion(mult = c(0, 0.02)),
                       labels = scales::percent_format(accuracy = 1)) +
    theme_ppmf() +
    expand_limits(x = 0) +
    expand_limits(x = 2) +
    labs(x = glue::glue("VI Distance"),
         y = NULL,
         caption = glue("Sample of {length(use_ind)} plans from {formatC(n_distinct(obj$draw), big.mark = ',')} draws"),
         title = title) +
    theme(plot.title = ggtext::element_markdown())
}


# boxplots ----


#' @param geo the value of `state` to subset to
#' @param show the vector of `source` plans to show
#' @param grouped NULL if districts can be treated as is or a length-two named vector like
#'  c(by = 4, max = 124) where [by] indicates the group length and [max] is the cap.
#'
plot_boxes = function(data, geo, title, grouped = NULL, show) {

  data <- data %>%
    filter(state == geo)

  if (!is.null(grouped)) {
    data <- data %>%
      mutate(district = fct_inorder(str_c(district, "-", pmin(district + grouped["by"], grouped["max"])), ordered=T))
  }

  data %>%
    mutate(source = lbl_source(source)) %>%
    filter(source %in% show) %>%
    ggplot(aes(x = as.factor(district),
               fill = source,
               ymin = low,
               lower = q1,
               middle = med,
               upper = q3,
               ymax = high)) +
    geom_hline(yintercept=0, lty="dashed", color="#444444") +
    geom_boxplot(width=0.9, size=0.25, stat = "identity") +
    scale_fill_manual(values=PAL_DAS) +
    scale_y_continuous("Difference from enacted plan",
                       labels = percent_format(accuracy=1, suffix="pp")) +
    labs(x=NULL, title = title, fill = "Plans sampled from") +
    theme_ppmf()
}



